#!/bin/bash

# Copyright 2017 Xingyu Na
# Apache 2.0

. ./path.sh || exit 1;



# 新增：子集比例参数
subset_ratio=0.20  # 默认10%

if [ "$1" == "--subset-ratio" ]; then
    subset_ratio=$2
    shift 2
fi

if [ $# != 2 ]; then
  echo "Usage: $0 [--subset-ratio 0.1] <audio-path> <text-path>"
  echo " $0 /export/a05/xna/data/data_aishell/wav /export/a05/xna/data/data_aishell/transcript"
  exit 1;
fi

aishell_audio_dir=$1
aishell_text=$2/aishell_transcript_v0.8.txt

train_dir=data/local/train
dev_dir=data/local/dev
test_dir=data/local/test
tmp_dir=data/local/tmp

# train_dir=/root/autodl-tmp/data_aishell/local/train
# dev_dir=/root/autodl-tmp/data_aishell/local/dev
# test_dir=/root/autodl-tmp/data_aishell/local/test
# tmp_dir=/root/autodl-tmp/data_aishell/local/tmp

mkdir -p $train_dir
mkdir -p $dev_dir
mkdir -p $test_dir
mkdir -p $tmp_dir

# data directory check
if [ ! -d $aishell_audio_dir ] || [ ! -f $aishell_text ]; then
  echo "Error: $0 requires two directory arguments"
  exit 1;
fi

# find wav audio file for train, dev and test resp.
find $aishell_audio_dir -iname "*.wav" > $tmp_dir/wav.flist
n=`cat $tmp_dir/wav.flist | wc -l`
[ $n -ne 141925 ] && \
  echo Warning: expected 141925 data data files, found $n

# 新增：按说话人ID抽取子集
mkdir -p $tmp_dir/speaker_subsets

# 从文件名中提取说话人ID（格式为 BAC009SXXXXWYYYY）
for wav in $(cat $tmp_dir/wav.flist); do
    filename=$(basename $wav)
    speaker=$(echo $filename | grep -oP 'S\d+')  # 提取 S 后面跟着数字的部分
    echo $wav >> $tmp_dir/speaker_subsets/${speaker}.list
done


# 为每个说话人随机选择指定比例的音频
for speaker_file in $tmp_dir/speaker_subsets/*.list; do
    speaker=$(basename $speaker_file .list)
    total_samples=$(wc -l < $tmp_dir/speaker_subsets/${speaker}.list)
    subset_size=$((total_samples * 1 / 100))  # 1%，使用整数运算
    [ $subset_size -lt 1 ] && subset_size=1   # 至少保留1个样本
  
    # 随机选择样本
    shuf $speaker_file | head -n $subset_size >> $tmp_dir/subset_wav.flist
done

# 按原始逻辑分类子集
grep -i "/train/" $tmp_dir/subset_wav.flist > $train_dir/wav.flist
grep -i "/dev/" $tmp_dir/subset_wav.flist > $dev_dir/wav.flist
grep -i "/test/" $tmp_dir/subset_wav.flist > $test_dir/wav.flist

rm -r $tmp_dir

# Transcriptions preparation
for dir in $train_dir $dev_dir $test_dir; do
  echo Preparing $dir transcriptions
  sed -e 's/\.wav//' $dir/wav.flist | awk -F '/' '{print $NF}' > $dir/utt.list
  paste -d' ' $dir/utt.list $dir/wav.flist > $dir/wav.scp_all
  tools/filter_scp.pl -f 1 $dir/utt.list $aishell_text > $dir/transcripts.txt
  awk '{print $1}' $dir/transcripts.txt > $dir/utt.list
  tools/filter_scp.pl -f 1 $dir/utt.list $dir/wav.scp_all | sort -u > $dir/wav.scp
  sort -u $dir/transcripts.txt > $dir/text
done

mkdir -p data/train data/dev data/test

# mkdir -p /root/autodl-tmp/data_aishell/train
# mkdir -p /root/autodl-tmp/data_aishell/dev
# mkdir -p /root/autodl-tmp/data_aishell/test

for f in wav.scp text; do
  cp $train_dir/$f data/train/$f || exit 1;
  cp $dev_dir/$f data/dev/$f || exit 1;
  cp $test_dir/$f data/test/$f || exit 1;
done

# for f in wav.scp text; do
#   cp $train_dir/$f /root/autodl-tmp/data_aishell/train/$f || exit 1;
#   cp $dev_dir/$f /root/autodl-tmp/data_aishell/dev/$f || exit 1;
#   cp $test_dir/$f /root/autodl-tmp/data_aishell/test/$f || exit 1;
# done

echo "$0: AISHELL data preparation succeeded"
exit 0;